#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
测试使用通义千问模型进行中文JSON键名映射
"""

import os
import json
import time
import tabulate
from litellm_openai_compat.map_qwen_chinese_json import map_chinese_to_english, key_mapping_results

def print_mapping_summary(mapping_results):
    """
    打印中英文键名映射摘要
    
    Args:
        mapping_results: 键名映射结果字典
    """
    print("\n" + "="*60)
    print("  中英文键名映射摘要")
    print("="*60)
    
    table_data = []
    for key, mapped_keys in mapping_results.items():
        if len(mapped_keys) == 1:
            table_data.append([key, mapped_keys[0]])
        else:
            table_data.append([key, "\n".join(mapped_keys)])
    
    # 使用tabulate打印表格
    try:
        print(tabulate.tabulate(
            table_data, 
            headers=["中文键名", "英文键名"], 
            tablefmt="grid"
        ))
    except ImportError:
        # 如果没有tabulate模块，使用简单格式
        print("中文键名\t英文键名")
        print("-"*40)
        for row in table_data:
            print(f"{row[0]}\t{row[1]}")
    
    print("="*60)

def print_json_sample(data, max_items=5):
    """
    打印JSON数据样例
    
    Args:
        data: JSON数据
        max_items: 最大显示项数
    """
    if isinstance(data, dict):
        sample = {}
        for i, (key, value) in enumerate(data.items()):
            if i >= max_items:
                sample["..."] = "..."
                break
            if isinstance(value, (dict, list)):
                sample[key] = "(complex structure)"
            else:
                sample[key] = value
        print(json.dumps(sample, ensure_ascii=False, indent=2))
    else:
        print(json.dumps(data, ensure_ascii=False, indent=2))

def main():
    # 定义输入、映射和输出文件路径
    input_file = "litellm_openai_compat/invoice_files/rec.json"
    mapping_file = "litellm_openai_compat/invoice_files/target_c2e_mapping.json"
    output_file = "litellm_openai_compat/invoice_files/rec_english.json"
    
    print("\n" + "="*60)
    print("  中文JSON键名映射工具测试")
    print("="*60)
    print(f"输入文件: {input_file}")
    print(f"映射文件: {mapping_file}")
    print(f"输出文件: {output_file}")
    
    # 加载原始JSON数据
    with open(input_file, "r", encoding="utf-8") as f:
        original_data = json.load(f)
    
    print("\n[原始JSON数据示例]")
    print_json_sample(original_data)
    
    print("\n开始进行中文JSON键名映射...")
    start_time = time.time()
    
    # 调用映射函数
    result = map_chinese_to_english(input_file, mapping_file, output_file)
    
    elapsed_time = time.time() - start_time
    
    if result:
        print(f"\n映射成功完成！耗时: {elapsed_time:.2f}秒")
        print("\n[转换后的JSON数据示例]")
        print_json_sample(result)
        
        # 打印键名映射摘要
        print_mapping_summary(key_mapping_results)
        
        # 统计映射情况
        total_keys = len(key_mapping_results)
        split_keys = sum(1 for keys in key_mapping_results.values() if len(keys) > 1)
        direct_mapped = sum(1 for key, keys in key_mapping_results.items() 
                           if len(keys) == 1 and key in key_mapping_results)
        
        print(f"\n总计转换 {total_keys} 个键名")
        print(f"其中拆分为多个键的复合字段: {split_keys} 个")
        print(f"直接映射的字段: {direct_mapped} 个")
        print(f"通过语义匹配的字段: {total_keys - split_keys - direct_mapped} 个")
        
        print(f"\n转换后的完整JSON已保存到: {output_file}")
        print(f"键名映射结果已保存到: {os.path.splitext(output_file)[0]}_mapping_result.json")
    else:
        print("映射过程中发生错误。")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n程序被用户中断")
    except Exception as e:
        print(f"发生错误: {e}") 